import pandas as pd
import os

# 当前目录
current_dir = os.path.dirname(os.path.abspath(__file__))

# 设置显示选项
pd.set_option('display.max_columns', None)
pd.set_option('display.width', 1000)

# 读取数据文件
def load_data():
    print("正在加载数据...")
    # 注意：除了word_translation.csv外，其他文件使用'>'作为分隔符
    book_df = pd.read_csv(os.path.join(current_dir, 'book.csv'), sep=">")
    relation_df = pd.read_csv(os.path.join(current_dir, 'relation_book_word.csv'), sep=">")
    word_df = pd.read_csv(os.path.join(current_dir, 'word.csv'), sep=">")
    word_translation_df = pd.read_csv(os.path.join(current_dir, 'word_translation.csv'))
    
    print("数据加载完成!")
    return book_df, relation_df, word_df, word_translation_df

# 查找指定的单词书
def find_book(book_df, book_name):
    # 在书名中搜索
    matching_books = book_df[book_df['bk_name'].str.contains(book_name, na=False)]
    
    if matching_books.empty:
        # 尝试在完整书名中搜索
        matching_books = book_df[book_df['bk_book'].str.contains(book_name, na=False)]
    
    return matching_books

# 提取指定单词书中的单词
def extract_book_words(book_name):
    # 加载数据
    book_df, relation_df, word_df, word_translation_df = load_data()
    
    # 查找单词书
    print(f"\n正在查找单词书: '{book_name}'")
    matching_books = find_book(book_df, book_name)
    
    if matching_books.empty:
        print(f"未找到名为 '{book_name}' 的单词书")
        # 显示一些包含"外研版"的单词书，供用户参考
        wy_books = book_df[book_df['bk_name'].str.contains("外研版", na=False)]
        if not wy_books.empty:
            print("\n找到以下包含'外研版'的单词书：")
            for idx, row in wy_books.iterrows():
                print(f"- {row['bk_name']}")
        return
    
    # 如果找到多本匹配的书，让用户选择
    if len(matching_books) > 1:
        print(f"\n找到{len(matching_books)}本匹配的单词书：")
        for i, (idx, row) in enumerate(matching_books.iterrows()):
            print(f"{i+1}. {row['bk_name']} (单词数: {row['bk_item_num']})")
        
        # 由于是脚本运行，默认选择第一本
        selected_index = 0
        selected_book = matching_books.iloc[selected_index]
        print(f"\n选择第{selected_index+1}本: {selected_book['bk_name']}")
    else:
        selected_book = matching_books.iloc[0]
        print(f"找到单词书: {selected_book['bk_name']}")
        print(f"- 作者: {selected_book['bk_author']}")
        print(f"- 出版社: {selected_book['bk_publisher']}")
        print(f"- 单词总数: {selected_book['bk_item_num']}")
    
    # 获取单词书ID
    book_id = selected_book['bk_id']
    
    # 查找该单词书中的所有单词关系
    book_relations = relation_df[relation_df['bv_book_id'] == book_id]
    
    if book_relations.empty:
        print("未找到该单词书中的单词关系")
        return
    
    # 合并单词信息
    merged_words = pd.merge(book_relations, word_df, left_on='bv_voc_id', right_on='vc_id')
    
    # 合并翻译信息
    final_result = pd.merge(merged_words, word_translation_df, left_on='vc_vocabulary', right_on='word', how='left')
    
    # 选择需要的列并排序
    columns_to_keep = [
        'bv_order', 'vc_vocabulary', 'vc_phonetic_uk', 'vc_phonetic_us', 
        'translation', 'bv_tag', 'vc_frequency', 'vc_difficulty', 'vc_acknowledge_rate'
    ]
    
    result_sorted = final_result[columns_to_keep].sort_values('bv_order')
    
    # 重命名列，使其更直观
    result_sorted.rename(columns={
        'bv_order': '序号',
        'vc_vocabulary': '单词',
        'vc_phonetic_uk': '英音音标',
        'vc_phonetic_us': '美音音标',
        'translation': '翻译',
        'bv_tag': '单元',
        'vc_frequency': '词频',
        'vc_difficulty': '难度',
        'vc_acknowledge_rate': '认识率'
    }, inplace=True)
    
    # 保存结果到CSV文件
    output_file = os.path.join(current_dir, f"{book_name.replace('/', '_').replace('\\', '_')}_单词表.csv")
    result_sorted.to_csv(output_file, index=False, encoding='utf-8-sig')
    
    print(f"\n已成功提取{len(result_sorted)}个单词到文件: {output_file}")
    
    # 显示前10个单词作为预览
    print("\n单词表预览 (前10个):")
    print(result_sorted[['序号', '单词', '美音音标', '翻译', '单元']].head(10))

if __name__ == "__main__":
    # 要提取的单词书名称
    target_book = "外研版三年级起点四年级上"
    extract_book_words(target_book)